package ai;

import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.util.EntityUtils;
import org.json.JSONArray;
import org.json.JSONObject;

public class OllamaDeepSeekR1Demo {

    // Ollama 服务地址（确认服务正在运行）
    private static final String OLLAMA_API_URL = "http://localhost:11434/api/chat";
    // 实际模型名称根据安装情况调整（通过 ollama list 查看）
    private static final String MODEL_NAME = "deepseek-r1:1.5b";

    public static void main(String[] args) {
        String question = "量子计算机的基本工作原理是什么？";

        try {
            String response = askOllama(question);
            System.out.println("=== 完整响应 ===");
            System.out.println(response);

            String answer = parseResponse(response);
            System.out.println("\n=== 解析后的回答 ===");
            System.out.println(answer);
        } catch (Exception e) {
            System.err.println("请求出错: " + e.getMessage());
            e.printStackTrace();
        }
    }

    private static String askOllama(String prompt) throws Exception {
        try (CloseableHttpClient httpClient = HttpClients.createDefault()) {
            HttpPost httpPost = new HttpPost(OLLAMA_API_URL);

            // 构建符合 Ollama Chat API 规范的请求体
            JSONObject requestBody = new JSONObject();
            requestBody.put("model", MODEL_NAME);

            JSONArray messages = new JSONArray();
            messages.put(new JSONObject()
                    .put("role", "user")
                    .put("content", prompt));
            requestBody.put("messages", messages);

            requestBody.put("stream", false); // 关闭流式输出
            requestBody.put("temperature", 0.7); // 控制生成随机性

            // 设置请求头
            httpPost.setHeader("Content-Type", "application/json");
            httpPost.setEntity(new StringEntity(requestBody.toString(), "UTF-8"));

            // 发送请求并获取响应
            try (CloseableHttpResponse response = httpClient.execute(httpPost)) {
                int statusCode = response.getStatusLine().getStatusCode();
                if (statusCode != 200) {
                    String errorMsg = EntityUtils.toString(response.getEntity());
                    throw new RuntimeException("HTTP错误: " + statusCode + "\n" + errorMsg);
                }
                return EntityUtils.toString(response.getEntity(), "UTF-8");
            }
        }
    }

    private static String parseResponse(String jsonResponse) {
        JSONObject json = new JSONObject(jsonResponse);
        return json.getJSONObject("message").getString("content");
    }
}